# california_housing_regression.py
# 加州房价预测全流程：数据加载 → 预处理 → 模型训练 → 评估 → 可视化

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.pipeline import make_pipeline

# 设置随机种子保证可复现性
np.random.seed(42)


# --------------------- 1. 数据加载与探索 ---------------------
def load_data():
    """加载加州房价数据集并返回特征和标签"""
    data = fetch_california_housing()
    X = pd.DataFrame(data.data, columns=data.feature_names)
    y = pd.Series(data.target, name='MedHouseVal')

    print("\n=== 数据概览 ===")
    print("特征矩阵形状:", X.shape)
    print("标签形状:", y.shape)
    print("\n前5行特征数据:\n", X.head())
    print("\n特征描述:\n", X.describe())

    return X, y


# --------------------- 2. 数据预处理 ---------------------
def preprocess_data(X, y):
    """
    数据预处理：
    - 划分训练集/测试集
    - 特征标准化
    """
    # 划分训练集和测试集 (80%训练, 20%测试)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    # 特征标准化 (拟合训练集并转换训练集和测试集)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    print("\n=== 预处理结果 ===")
    print("训练集规模:", X_train_scaled.shape)
    print("测试集规模:", X_test_scaled.shape)

    return X_train_scaled, X_test_scaled, y_train, y_test


# --------------------- 3. 模型训练与评估 ---------------------
def train_and_evaluate(X_train, X_test, y_train, y_test):
    """训练多个模型并评估性能"""
    # 初始化模型
    models = {
        "Linear Regression": LinearRegression(),
        "Random Forest": RandomForestRegressor(n_estimators=100, random_state=42)
    }

    results = {}

    for name, model in models.items():
        # 训练模型
        model.fit(X_train, y_train)

        # 预测
        y_pred = model.predict(X_test)

        # 评估指标
        rmse = np.sqrt(mean_squared_error(y_test, y_pred))
        r2 = r2_score(y_test, y_pred)

        results[name] = {
            "RMSE": rmse,
            "R²": r2,
            "model": model
        }

        print(f"\n=== {name} ===")
        print(f"RMSE: {rmse:.4f}")
        print(f"R² Score: {r2:.4f}")

    return results


# --------------------- 4. 结果可视化 ---------------------
def visualize_results(y_test, y_pred, model_name):
    """绘制真实值 vs 预测值散点图"""
    plt.figure(figsize=(8, 6))
    plt.scatter(y_test, y_pred, alpha=0.5)
    plt.plot([y_test.min(), y_test.max()],
             [y_test.min(), y_test.max()],
             'r--', linewidth=2)
    plt.xlabel("True Values")
    plt.ylabel("Predictions")
    plt.title(f"{model_name}: True vs Predicted Values")
    plt.grid(True)
    plt.show()


def plot_feature_importance(model, feature_names):
    """绘制随机森林特征重要性"""
    if hasattr(model, 'feature_importances_'):
        importances = model.feature_importances_
        indices = np.argsort(importances)[::-1]

        plt.figure(figsize=(10, 6))
        plt.title("Feature Importances")
        plt.bar(range(len(indices)), importances[indices], align="center")
        plt.xticks(range(len(indices)), [feature_names[i] for i in indices], rotation=90)
        plt.xlim([-1, len(indices)])
        plt.tight_layout()
        plt.show()


# --------------------- 主程序 ---------------------
if __name__ == "__main__":
    # 1. 加载数据
    X, y = load_data()

    # 2. 数据预处理
    X_train, X_test, y_train, y_test = preprocess_data(X, y)

    # 3. 训练与评估模型
    results = train_and_evaluate(X_train, X_test, y_train, y_test)

    # 4. 可视化结果
    for name, result in results.items():
        # 绘制预测结果散点图
        y_pred = result["model"].predict(X_test)
        visualize_results(y_test, y_pred, name)

        # 如果是随机森林，绘制特征重要性
        if name == "Random Forest":
            plot_feature_importance(result["model"], X.columns)

    # 5. 输出最佳模型结果
    best_model = min(results.items(), key=lambda x: x[1]["RMSE"])
    print(f"\n*** 最佳模型: {best_model[0]} ***")
    print(f"RMSE: {best_model[1]['RMSE']:.4f}")
    print(f"R²: {best_model[1]['R²']:.4f}")